import matplotlib.pyplot as plt
import numpy as np

import matplotlib.animation as animation

fig, (heading, board) = plt.subplots(2, 1)

t = np.linspace(0, 3, 40)
g = -9.81
v0 = 12
z = g * t**2 / 2 + v0 * t

v02 = 5
z2 = g * t**2 / 2 + v02 * t

scat = board.scatter(t[0], z[0], c="b", s=5, label=f"v0 = {v0} m/s")
line2 = board.plot(t[0], z2[0], label=f"v0 = {v02} m/s")[0]
board.set(xlim=[0, 3], ylim=[-4, 10], xlabel="Time [s]", ylabel="Z [m]")
board.legend()


def update(frame):
    # for each frame, update the data stored on each artist.
    x = t[:frame]
    y = z[:frame]
    # update the scatter plot:
    data = np.stack([x, y]).T
    scat.set_offsets(data)
    # update the line plot:
    line2.set_xdata(t[:frame])
    line2.set_ydata(z2[:frame])
    return (scat, line2)


ani = animation.FuncAnimation(fig=fig, func=update, frames=40, interval=30)

# Try with different writers
try:
    ani.save("animation.gif", writer="pillow", fps=30)  # GIF format
    # or
    # ani.save("animation.webm", writer="ffmpeg", fps=30)  # WebM format
except Exception as e:
    print(f"Could not save animation: {e}")

plt.show()
